import math
import os.path

import numpy as np
import matplotlib.pyplot as plt


def get_data(res_path):
    name = ['lr', 'logit_scale', 'contrastive_loss', 'distill_loss', 'loss']
    res = {
        "lr": [],
        "logit_scale": [],
        "contrastive_loss": [],
        "distill_loss": [],
        "loss": [],
    }

    with open(res_path, 'r') as fp:
        for line in fp:
            if "Epoch" in line:
                data = line.strip()
                data = data.split('LR')[1].split(' ')
                name_index = 0
                for info_str in data:
                    try:
                        res[name[name_index]].append(float(info_str))
                        name_index += 1
                    except ValueError:
                        pass
                    except IndexError:
                        break
    print(res)
    return res

def plt_figure(save_path, data):
    assert isinstance(data, list)

    for i in range(0, len(data)):
        if not math.isnan(data[i]):
            plt.scatter(i, data[i], color='blue', s=2)
        else:
            y = plt.ylim()[0] * 15
            plt.scatter(i, y, color='red', s=2)

    plt.xlabel('Epoch')
    plt.ylabel(save_path.split('\\')[-1].split('.')[0])

    plt.savefig(save_path)
    plt.close()


if __name__ == '__main__':
    path_res = r"D:\Project\multimoding\res\CLIP_Pretrain\3\out.log"
    path_save = r"D:\Project\multimoding\res\CLIP_Pretrain\3"
    res = get_data(path_res)

    for name, data in res.items():
        plt_figure(os.path.join(path_save, ''.join([name, '.png'])), data)
